// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/runtime/container_private.h>
#include <tbir/runtime/registry.h>

namespace {

    using namespace ::tbir::runtime;

// only an simple example
    static NDArray tensor_ops_zeros(const RTValue &shape) {
        if (!shape.IsObjectRef<List>()) {
            MXTHROW << "expect shape is list type, but get " << shape.type_name();
        }
        int64_t bytes_len = 1;
        auto ls = shape.AsObjectRefNoCheck<List>();
        std::vector<int64_t> ft_shape;
        for (const auto &di : ls) {
            ft_shape.push_back(di.As<int64_t>());
            bytes_len *= ft_shape.back();
        }
        if (ft_shape.empty()) {
            bytes_len = 0;
        }
        auto ret = NDArray::Empty(ft_shape, {kDLFloat, 64, 1}, {kDLCPU, 0});
        auto dptr = (char *) (ret->data) + ret->byte_offset;
        memset(dptr, 0, bytes_len);
        return ret;
    }

    TBIR_REGISTER_NATIVE_NAMED_FUNC("numpy_ops_zeros", tensor_ops_zeros);
    TBIR_REGISTER_NATIVE_NAMED_FUNC("torch_ops_zeros", tensor_ops_zeros);

}  // namespace
